import base64
import socket
import uuid
from Ios import Client, Device, Usbmux
from multiprocessing import Process
from Common import *
from Driver.BaseDriver import BaseDriver, check_status
from typing import Union


class IosDriver(BaseDriver):

    def __init__(self, serial: str, port: int = None):
        self._process: WdaProcess = None
        self._client: Client = None
        self._serial = serial
        self._platform = 'ios'
        self._port = port
        self.status = 'disconnect'

    def _get_not_use_port(self) -> int:
        """获取未使用的端口，若初始化已经携带端口，直接返回该端口"""
        if self._port is not None:
            return self._port
        port = 8100
        while True:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                s.connect(('127.0.0.1', port))
                s.shutdown(2)
                self._port = port
                return
            except:
                port += 1

    ###########################################################################
    # 设备方法
    ###########################################################################
    def device_connect(self):
        """
        连接设备，并启动wda
        :return:
        """
        self._get_not_use_port()
        self._client = Client("http://127.0.0.1:{}".format(self._port))
        if self._client.is_ready():
            return
        self._process = WdaProcess(self._port)
        if self._client.wait_ready() is False:
            raise ConnectError("[http://127.0.0.1:{}] connect fail".format(self._port))
        self._client.home()
        self._client.healthcheck()
        self.status = 'connect'

    def device_disconnect(self):
        """
        断开连接
        :return:
        """
        if self._process:
            self._process.terminate()
            self._process.join()
        self._client.close()
        self.status = 'disconnect'

    ###########################################################################
    # driver方法
    ###########################################################################
    @check_status
    def session_click(self, x: float, y: float):
        """
        点击事件
        :param x: x坐标
        :param y: y坐标
        :return:
        """
        return self._client.click(int(x/self._client.scale), int(y/self._client.scale))

    @check_status
    def session_swipe(self, fx: int, fy: int, tx: int, ty: int):
        """
        滑动事件
        :param fx: 起点x坐标
        :param fy: 起点y坐标
        :param tx: 终点x坐标
        :param ty: 终点y坐标
        :return:
        """
        return self._client.swipe(int(fx / self._client.scale), int(fy / self._client.scale),
                                  int(tx / self._client.scale), int(ty / self._client.scale))

    @check_status
    def session_system(self, key: Union[int, str]):
        """
        系统按键事件
        :param key: 按键名
        :return:
        """
        if key in ("home", "volumeUp", "volumeDown"):
            return self._client.press(key)
        elif key == "enter":
            return self._client.send_keys('\r')
        raise RuntimeError('key not support')

    @check_status
    def session_info(self):
        """
        系统信息
        :return:
        """
        return self._client.info

    @check_status
    def session_screenshot(self):
        """
        截图
        :return:
        """
        return str(base64.b64encode(self._client.screenshot(format="raw")).decode('ascii'))

    @check_status
    def session_dump(self):
        """
        获取xml控件树
        :return:
        """
        return self._get_ios_hierarchy(self._client)

    @check_status
    def session_windows(self):
        """
        获取设备显示大小
        :return:
        """
        return self._client.window_size()

    def _get_ios_hierarchy(self, d):
        """
        获取控件树
        :param d: client
        :return:
        """
        sourcejson = d.source(format='json')

        def travel(node, parent):
            if parent:
                node['_parent'] = parent
            node['_id'] = str(uuid.uuid4())
            node['_type'] = node.pop('type', "null")
            if node.get('rect'):
                rect = node['rect']
                nrect = {}
                for k, v in rect.items():
                    nrect[k] = v * d.scale
                node['rect'] = nrect
            node['accessibility id'] = node['label']
            node['label'] = 'XCUIElementType' + node['_type']
            for child in node.get('children', []):
                travel(child, node['_id'])
            return node

        return travel(sourcejson, None)

    ###########################################################################
    # driver方法
    ###########################################################################
    @check_status
    def app_start(self, app: str):
        """
        启动app
        :param app: app包名
        :return:
        """
        return self._client.app_start(app)

    @check_status
    def app_stop(self, app: str):
        """
        停止app
        :param app: app包名
        :return:
        """
        return self._client.app_stop(app)

    @check_status
    def app_restart(self, app: str):
        """
        重启app
        :param app: app包名
        :return:
        """
        self._client.app_stop(app)
        return self._client.app_start(app)

    @check_status
    def app_reset(self, app: str):
        """
        重置app
        :param app: app包名
        :return:
        """
        pass

    ###########################################################################
    # element方法
    ###########################################################################
    @check_status
    def element_click(self, locator):
        """
        点击控件
        :param locator: 控件定位器
        :return:
        """
        locator = IosLocator.create_locator(locator)
        if locator.locator_offset == (0, 0):
            return self._client(**locator.locator_arg).click(5) if locator.locator_index == -1 \
                else self._client(**locator.locator_arg)[locator.locator_index].click(5)
        else:
            element_info = self.element_info(locator)
            click_x = (element_info['bounds']['left'] + element_info['bounds']['right']) / 2 + locator.locator_offset[0]
            click_y = (element_info['bounds']['top'] + element_info['bounds']['bottom']) / 2 + locator.locator_offset[1]
            return self.session_click(click_x, click_y)

    @check_status
    def element_send(self, locator, text: str):
        """
        向控件发送文字
        :param locator: 控件定位器
        :param text: 文本
        :return:
        """
        locator = IosLocator.create_locator(locator)
        return self._client(**locator.locator_arg).set_text(text) if locator.locator_index == -1 \
            else self._client(**locator.locator_arg)[locator.locator_index].set_text(text)

    @check_status
    def element_clear(self, locator):
        """
        清空控件文本
        :param locator: 控件定位器
        :return:
        """
        locator = IosLocator.create_locator(locator)
        return self._client(**locator.locator_arg).clear_text(5) if locator.locator_index == -1 \
            else self._client(**locator.locator_arg)[locator.locator_index].clear_text(5)

    @check_status
    def element_exists(self, locator):
        """
        判断控件是否存在
        :param locator: 控件定位器
        :return:
        """
        locator = IosLocator.create_locator(locator)
        return self._client(**locator.locator_arg).exists() if locator.locator_index == -1 \
            else self._client(**locator.locator_arg)[locator.locator_index].exists()

    @check_status
    def element_info(self, locator):
        """
        获取控件信息
        :param locator: 控件定位器
        :return:
        """
        locator = IosLocator.create_locator(locator)
        ret = self._client(**locator.locator_arg).info if locator.locator_index == -1 \
            else self._client(**locator.locator_arg)[locator.locator_index].info
        rect = self._client(**locator.locator_arg).bounds if locator.locator_index == -1 \
            else self._client(**locator.locator_arg)[locator.locator_index].bounds
        ret['bounds'] = {'left': rect.left, 'right': rect.right, 'top': rect.top, 'bottom': rect.bottom}
        return ret

    ###########################################################################
    # find方法
    ###########################################################################
    @check_status
    def find_element(self, locator):
        """
        查找控件
        :param locator: 控件定位器
        :return:
        """
        return self.element_exists(locator)

    @check_status
    def find_elements(self, locator):
        """
        查找控件（多个）
        :param locator: 控件定位器
        :return:
        """
        locator = IosLocator.create_locator(locator)
        return len(self._client.xpath(locator.locator_value).all()) if locator.locator_type == 'xpath' \
            else self._client(**locator.locator_arg).count

    @check_status
    def find_allinfo(self, locator):
        element_count = self.find_elements(locator)
        result_list = []
        for i in range(0, element_count):
            locator['index'] = i
            info = self.element_info(locator)
            if 'visibleBounds' in info:
                result_list.append({'x': info['visibleBounds']['left'], 'y': info['visibleBounds']['top'],
                                    'width': info['visibleBounds']['right'] - info['visibleBounds']['left'],
                                    'height': info['visibleBounds']['bottom'] - info['visibleBounds']['top']})
            else:
                result_list.append({'x': info['bounds']['left'], 'y': info['bounds']['top'],
                                    'width': info['bounds']['right'] - info['bounds']['left'],
                                    'height': info['bounds']['bottom'] - info['bounds']['top']})
        return result_list

    @classmethod
    def get_devices(cls):
        devices = []
        for device in Usbmux().device_list():
            devices.append(device['UDID'])
        return devices


class WdaProcess(Process):

    def __init__(self, d: str, port: int) -> None:
        super().__init__()
        self._d = d
        self._port = port

    @property
    def port(self):
        """获取设备占用的本机端口"""
        return self._port

    @property
    def udid(self):
        """获取进程设备udid"""
        return self._d

    def run(self) -> None:
        """
        启动wda
        :return:
        """
        d = Device(self._d)
        d.xctest(self.port, parent=self)


if __name__ == '__main__':
    # print(IosDriver("1"))
    print(BaseDriver.get_method_names('device'))
    # for n in BaseDriver.get_method_names():
    #     print(n)
    #     print(BaseDriver.get_method_args(n))